/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 *  The ASF licenses this file to You under the Apache License, Version 2.0
 *  (the "License"); you may not use this file except in compliance with
 *  the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 *  Unless required by applicable law or agreed to in writing, software
 *  distributed under the License is distributed on an "AS IS" BASIS,
 *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *  See the License for the specific language governing permissions and
 *  limitations under the License.
 */

package com.example.ai.agent.controller;

import com.example.ai.agent.mcp.tool.MathCalculator;
import com.example.ai.agent.mcp.tool.WeatherService;
import jakarta.annotation.PostConstruct;
import jakarta.servlet.http.HttpServletResponse;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import reactor.core.publisher.Flux;

import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.SimpleLoggerAdvisor;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.ollama.api.OllamaOptions;
import org.springframework.ai.model.function.FunctionCallback;

import java.util.function.Function;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.core.env.Environment;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;

/**
 * @author yuluo
 * @author <a href="mailto:yuluo08290126@gmail.com">yuluo</a>
 */

@RestController
@RequestMapping("/client")
public class OllamaChatClientController {

	private static final Logger log = LoggerFactory.getLogger(OllamaChatClientController.class);

	@Value("${spring.ai.ollama.chat.model}")
	private String model;

	@Value("${spring.ai.ollama.base-url}")
	private String baseUrl;

	private ChatClient ollamaiChatClient;

	@Autowired
	private Environment environment;
	
	@Autowired
	private ChatModel chatModel;
	
	@Autowired
	private MathCalculator mathCalculator;
	
	@Autowired
	private WeatherService weatherService;

	@PostConstruct
	public void init() {
		// 在属性注入完成后，构建ChatClient，设置正确的模型参数
		this.ollamaiChatClient = ChatClient.builder(chatModel)
				// 实现 Logger 的 Advisor
				.defaultAdvisors(
						new SimpleLoggerAdvisor()
				)
				// 设置 ChatClient 中 ChatModel 的 Options 参数
				.defaultOptions(
						OllamaOptions.builder()
								.topP(0.7)
								.model(model)
								.build()
				)
				// 注册工具函数
				.defaultFunctions(
						FunctionCallback.builder()
						.function("calculate", mathCalculator::calculate)
						.inputType(MathCalculator.CalculateRequest.class)
						.description("执行数学计算，支持加、减、乘、除四则运算")
						.build(),
						FunctionCallback.builder()
						.function("getWeather", weatherService::getWeather)
						.inputType(WeatherService.WeatherRequest.class)
						.description("查询指定城市的天气信息，包括温度、天气状况、风向、风速和湿度")
						.build()
				)
				.build();
		
		log.info("OllamaChatClient 初始化完成，已注册数学计算和天气查询工具，使用模型: {}", model);
	}

	/**
	 * ChatClient 简单调用
	 */
	@GetMapping("/simple/chat")
	public String simpleChat(@RequestParam(value = "prompt", defaultValue = "你好，介绍下你自己！请用中文回答。") String prompt) {
		try {
			log.info("收到聊天请求，prompt: {}", prompt);
			String response = ollamaiChatClient.prompt(prompt).call().content();
			log.info("聊天响应成功，response长度: {}", response.length());
			return response;
		} catch (Exception e) {
			log.error("聊天请求处理失败", e);
			return "处理请求时发生错误: " + e.getMessage();
		}
	}

	/**
	 * ChatClient 流式调用
	 */
	@GetMapping("/stream/chat")
	public Flux<String> streamChat(@RequestParam(value = "prompt", defaultValue = "你好，介绍下你自己！请用中文回答。") String prompt, HttpServletResponse response) {
		try {
			log.info("收到流式聊天请求，prompt: {}", prompt);
			response.setCharacterEncoding("UTF-8");
			return ollamaiChatClient.prompt(prompt).stream().content();
		} catch (Exception e) {
			log.error("流式聊天请求处理失败", e);
			return Flux.just("处理请求时发生错误: " + e.getMessage());
		}
	}

	/**
	 * 获取当前配置信息
	 */
	@GetMapping("/config")
	public String getConfig() {
		String activeProfile = environment.getActiveProfiles().length > 0 ? environment.getActiveProfiles()[0] : "default";
		return String.format("当前使用的配置: %s, Ollama服务地址: %s, 模型: %s", activeProfile, baseUrl, model);
	}

}